import pandas as pd
from prophet import Prophet
from app.data.data_provider import get_stock_data
from datetime import timedelta

def run_forecast(code: str, days_to_predict: int = 30):
    """
    使用Prophet模型进行时间序列预测的核心函数。
    """
    # 1. 获取足够长的历史数据用于训练，比如过去3年
    end_date = pd.to_datetime("today")
    start_date = end_date - timedelta(days=3 * 365)
    
    df = get_stock_data(code, start_date.strftime('%Y-%m-%d'), end_date.strftime('%Y-%m-%d'))
    
    if df.empty or len(df) < 2:
        return {"error": "获取的历史数据不足以进行预测。"}

    # 2. 准备数据，列名必须是 'ds' 和 'y'
    df_prophet = df[['日期', '收盘']].copy()
    df_prophet.rename(columns={'日期': 'ds', '收盘': 'y'}, inplace=True)
    df_prophet['ds'] = pd.to_datetime(df_prophet['ds'])

    # --- 日志 ---
    print("\n--- [Forecast Engine] Prophet Input Data (tail) ---")
    print(df_prophet.tail())
    print(f"Input data length: {len(df_prophet)}")
    print("--------------------------------------------------\n")

    # 3. 初始化并训练模型，注入中国节假日信息
    model = Prophet(changepoint_prior_scale=0.05)
    model.add_country_holidays(country_name='CN') # <--- 关键改动 1
    model.fit(df_prophet)

    # 4. 创建未来的日期框架，频率为'B'(Business day)以排除周末
    future = model.make_future_dataframe(periods=days_to_predict, freq='B') # <--- 关键改动 2
    
    # 5. 执行预测
    forecast = model.predict(future)

    # --- 日志 ---
    print("\n--- [Forecast Engine] Prophet Forecast Output (tail) ---")
    print(forecast[['ds', 'yhat', 'yhat_lower', 'yhat_upper']].tail())
    print("------------------------------------------------------\n")

    # 6. 清理并返回结果
    forecast_data = forecast[['ds', 'yhat', 'yhat_lower', 'yhat_upper']].copy()
    forecast_data['ds'] = forecast_data['ds'].dt.strftime('%Y-%m-%d')
    df_prophet['ds'] = df_prophet['ds'].dt.strftime('%Y-%m-%d')
    
    return {
        "historical_data": df_prophet.to_dict('records'),
        "forecast_data": forecast_data.to_dict('records')
    } 